package org.luxor.commons.core.utils.smcrypto;

import cn.hutool.crypto.KeyUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.SM2;
import org.apache.commons.codec.binary.Base64;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.signers.SM2Signer;
import org.bouncycastle.jcajce.provider.asymmetric.util.ECUtil;

import java.io.Serializable;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;

/**
 * SM2国密算法类（对标RSA-非对称加、解密算法）
 *
 * @author Mr.yan  @date 2022/11/15
 */
public class Sm2Utils {
    private static final String ALGORITHM_SM2 = "SM2";
    public static final int KEY_SIZE = 1024;

    public static final String DEFAULT_PUBLIC_KEY_STR = "MFkwEwYHKoZIzj0CAQYIKoEcz1UBgi0DQgAElBTZ1scDX4+QQty5fvgCMi/Rux" +
            "ufJehF8iZZDmz31FiLq78H1NZ/KOIwin+L3IdrGrmy78GSfma6AkVg7DSQtQ==";

    public static final String DEFAULT_PRIVATE_KEY_STR = "MIGTAgEAMBMGByqGSM49AgEGCCqBHM9VAYItBHkwdwIBAQQgGGynZUcCvbMKm" +
            "0HKy6nCUNIbnl+E0TMtsr++CrpER9CgCgYIKoEcz1UBgi2hRANCAASUFNnWxwNfj5BC3Ll++AIyL9G7G58l6EXyJlkObPfUWIurvwfU1n8" +
            "o4jCKf4vch2saubLvwZJ+ZroCRWDsNJC1";

    public static void main(String[] args) {

        // 生成 私钥 & 公钥
        Keypair keypair = generateKeyPair();
        System.out.println("公钥: " + keypair.getPublicKey());
        System.out.println("私钥: " + keypair.getPrivateKey());

        String plainTextStr = "SM2国密算法类";
        System.out.println("明文内容: " + plainTextStr);

        System.out.println("--------------------");

        byte[] encryptText = encryptByPublicKey(plainTextStr.getBytes(), keypair.getPublicKey());
        System.out.println("1.1 公钥加密结果: " + Base64.encodeBase64String(encryptText));
        byte[] plainText = decryptByPrivateKey(encryptText, keypair.getPrivateKey());
        System.out.println("1.2 私钥解密结果: " + new String(plainText));

        System.out.println("--------------------");

        // 数字签名验证
        byte[] sign = sign(plainTextStr.getBytes(), DEFAULT_PRIVATE_KEY_STR);
        System.out.println("数字签名验证 verify:" + verify(plainTextStr.getBytes(), sign, DEFAULT_PUBLIC_KEY_STR));
    }


    /**
     * SM2公钥加密
     *
     * @param plaintext    明文
     * @param publicKeyStr 公钥串
     * @return 密文字节
     */
    public static byte[] encryptByPublicKey(byte[] plaintext, String publicKeyStr) {
        SM2 sm2 = SmUtil.sm2();
        byte[] keyBytes = Base64.decodeBase64(publicKeyStr.getBytes());
        sm2.setPublicKey(KeyUtil.generatePublicKey(ALGORITHM_SM2, keyBytes));
        return sm2.encrypt(plaintext, KeyType.PublicKey);
    }

    /**
     * SM2私钥解密
     *
     * @param ciphertext    密文
     * @param privateKeyStr 私钥串
     * @return java.lang.String
     */
    public static byte[] decryptByPrivateKey(byte[] ciphertext, String privateKeyStr) {
        SM2 sm2 = SmUtil.sm2();
        byte[] keyBytes = Base64.decodeBase64(privateKeyStr.getBytes());
        sm2.setPrivateKey(KeyUtil.generatePrivateKey(ALGORITHM_SM2, keyBytes));
        return sm2.decrypt(ciphertext, KeyType.PrivateKey);
    }


    /**
     * 依据私钥,生成数字签名
     *
     * @param data          数据
     * @param privateKeyStr 私钥串
     * @return 签名
     */
    public static byte[] sign(byte[] data, String privateKeyStr) {
        try {
            SM2Signer signer = new SM2Signer();
            byte[] keyBytes = Base64.decodeBase64(privateKeyStr.getBytes());
            PrivateKey privateKey = KeyUtil.generatePrivateKey(ALGORITHM_SM2, keyBytes);
            CipherParameters param = ECUtil.generatePrivateKeyParameter(privateKey);
            signer.init(true, param);
            signer.update(data, 0, data.length);
            return signer.generateSignature();
        } catch (Exception e) {
            throw new RuntimeException("生成数字签名失败", e);
        }
    }

    /**
     * 依据公钥,验证数字签名
     *
     * @param data         数据
     * @param sign         签名
     * @param publicKeyStr 公钥串
     * @return 验证结果
     */
    public static boolean verify(byte[] data, byte[] sign, String publicKeyStr) {
        try {
            SM2Signer signer = new SM2Signer();
            byte[] keyBytes = Base64.decodeBase64(publicKeyStr.getBytes());
            PublicKey publicKey = KeyUtil.generatePublicKey(ALGORITHM_SM2, keyBytes);
            CipherParameters param = ECUtil.generatePublicKeyParameter(publicKey);
            signer.init(false, param);
            signer.update(data, 0, data.length);
            return signer.verifySignature(sign);
        } catch (Exception e) {
            throw new RuntimeException("验证数字签名失败", e);
        }
    }

    /**
     * 生成一对 私钥 & 公钥
     *
     * @return KeyPair
     */
    public static Keypair generateKeyPair() {
        try {
            KeyPair pair = SecureUtil.generateKeyPair(ALGORITHM_SM2, KEY_SIZE);
            PrivateKey privateKey = pair.getPrivate();
            PublicKey publicKey = pair.getPublic();
            String privateKeyStr = Base64.encodeBase64String(privateKey.getEncoded());
            String publicKeyStr = Base64.encodeBase64String(publicKey.getEncoded());
            return new Sm2Utils.Keypair(privateKeyStr, publicKeyStr);
        } catch (Exception e) {
            throw new RuntimeException("生成SM2密钥对失败", e);
        }
    }

    static class Keypair implements Serializable {

        protected String privateKey;
        protected String publicKey;

        public Keypair() {
        }

        public Keypair(String privateKey, String publicKey) {
            this.privateKey = privateKey;
            this.publicKey = publicKey;
        }

        public String getPrivateKey() {
            return privateKey;
        }

        public void setPrivateKey(String privateKey) {
            this.privateKey = privateKey;
        }

        public String getPublicKey() {
            return publicKey;
        }

        public void setPublicKey(String publicKey) {
            this.publicKey = publicKey;
        }

        @Override
        public String toString() {
            return "Keypair{" +
                    "privateKey='" + privateKey + '\'' +
                    ", publicKey='" + publicKey + '\'' +
                    '}';
        }
    }
}
